package com.cy.sportassistant.agent.impl;

import com.cy.sportassistant.agent.SportAgent;
import com.cy.sportassistant.tools.BochaWebSearchEngine;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.memory.chat.TokenWindowChatMemory;
import dev.langchain4j.model.TokenCountEstimator;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiTokenCountEstimator;
import dev.langchain4j.rag.DefaultRetrievalAugmentor;
import dev.langchain4j.rag.RetrievalAugmentor;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.content.retriever.WebSearchContentRetriever;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.rag.query.router.DefaultQueryRouter;
import dev.langchain4j.rag.query.router.QueryRouter;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.milvus.MilvusEmbeddingStore;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.*;

import static dev.langchain4j.model.openai.OpenAiChatModelName.GPT_4_O_MINI;

@Slf4j
@Configuration
public class SportAgentImpl {

    @Bean
    ChatMemoryProvider chatMemoryProvider(TokenCountEstimator tokenizer) {
        return memoryId -> TokenWindowChatMemory.builder()
                .id(memoryId)
                .maxTokens(5000, tokenizer)
                .build();
    }


    @Bean
    EmbeddingStore<TextSegment> embeddingStore() {
        return MilvusEmbeddingStore.builder()
                .uri("https://in03-d1c03b1f6ba7563.serverless.ali-cn-hangzhou.cloud.zilliz.com.cn")
                .token("6dbc25361baceaa634d8271211479beaedcc63aac6e5bccdeef4cebd7a51dde59628dbfdd535b09ffad400c5ddcfa712d3736f96")
                .collectionName("sports_news")
                .dimension(1024)
                .build();
    }

    @Bean
    ContentRetriever contentRetriever(EmbeddingStore<TextSegment> embeddingStore, EmbeddingModel embeddingModel) {

        int maxResults = 5;
        double minScore = 0.7;

        return EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .maxResults(maxResults)
                .minScore(minScore)
                .build();
    }

    @Bean
    TokenCountEstimator tokenCountEstimator() {
        return new OpenAiTokenCountEstimator(GPT_4_O_MINI);
    }

    /**
     * 会话隔离、会话记忆、流式输出、functionCall
     *
     * @param streamingChatModel
     * @return
     */
    @Bean
    public SportAgent sportAgent(StreamingChatModel streamingChatModel, ContentRetriever embeddingStoreContentRetriever, ChatMemoryProvider chatMemoryProvider) {

        // Let's create our web search content retriever.
        //WebSearchEngine webSearchEngine = TavilyWebSearchEngine.builder()
        //        .baseUrl("https://api.bochaai.com/v1/web-search")
        //        //.apiKey("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") // get a free key: https://app.tavily.com/sign-in
        //        .apiKey("sk-ed5cd42588bb4863af67f68cd87f03f8") // get a free key: https://app.tavily.com/sign-in
        //        .build();

        BochaWebSearchEngine webSearchEngine = BochaWebSearchEngine.builder()
                .apiKey("sk-ed5cd42588bb4863af67f68cd87f03f8")
                //.freshness("noLimit")
                .freshness("oneWeek")
                .includeSummary(true)
                .build();

        ContentRetriever webSearchContentRetriever = WebSearchContentRetriever.builder()
                .webSearchEngine(webSearchEngine)
                .maxResults(3)
                .build();

        //QueryRouter queryRouter = new DefaultQueryRouter(embeddingStoreContentRetriever, webSearchContentRetriever);
        //QueryRouter queryRouter = new DefaultQueryRouter(embeddingStoreContentRetriever);

        // 创建动态QueryRouter

        //boolean enableWebSearch = (boolean) metadata.getOrDefault("enable_web_search", false);

        // 自定义QueryRouter
        QueryRouter queryRouter = new QueryRouter() {
            @Override
            public Collection<ContentRetriever> route(Query query) {
                String queryText = query.text();
                boolean needWebSearch = needWebSearch(queryText);

                log.info("Query: {}, Need Web Search: {}", queryText, needWebSearch);

                if (needWebSearch) {
                    log.info("Using both local knowledge base and web search");
                    return List.of(embeddingStoreContentRetriever, webSearchContentRetriever);
                } else {
                    log.info("Using only local knowledge base");
                    return List.of(embeddingStoreContentRetriever);
                }
            }
        };

        RetrievalAugmentor retrievalAugmentor = DefaultRetrievalAugmentor.builder()
                .queryRouter(queryRouter)
                .build();


        return AiServices.builder(SportAgent.class)
                //.tools(toolsService)
                //.toolProvider(toolProvider)
                .chatMemoryProvider(chatMemoryProvider)
                .retrievalAugmentor(retrievalAugmentor)
                .streamingChatModel(streamingChatModel)
                .build();
    }

    private boolean needWebSearch(String queryText) {
        // 定义不同类别的关键词
        Map<String, List<String>> keywordCategories = new HashMap<>();
        keywordCategories.put("时间相关", Arrays.asList("最新", "今天", "最近", "刚刚"));
        keywordCategories.put("新闻相关", Arrays.asList("新闻", "报道", "消息", "资讯"));
        keywordCategories.put("实时相关", Arrays.asList("实时", "直播", "现场", "即时"));

        // 检查是否包含任何类别的关键词
        return keywordCategories.values().stream()
                .flatMap(List::stream)
                .anyMatch(keyword -> queryText.toLowerCase().contains(keyword.toLowerCase()));
    }
}